import socket
import json
import threading
import time
class FLcloud:
    port=23334
    root=""
    root_state=None
    down_server=[]
    n_down_server=0
    down_server_state=[]

    sendjson1=json.dumps({"order":"aggregation_request"})
    sendjson2=json.dumps({"order":"test_local_model"})
    sendjson3=json.dumps({"order":"read_model"})
    sendjson4=json.dumps({"order":"save_model"})
    sendjson5=json.dumps({"order":"reset_model"})

    order={"1":"start a round of FL",
           "0":"exit"}
    running=True
    def __init__(self):
        return
    def listen(self):
        server_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        server_socket.bind(("", self.port)) 
        while(self.running):
            server_socket.settimeout(3)  #每3秒钟检测是否关闭
            try:
                receive_data, client = server_socket.recvfrom(1400)
                #print("来自客户端%s,发送的%s\n" % (client, receive_data))  #打印接收的内容
                json_data = json.loads(receive_data.decode('utf-8'))
                if (json_data['order']=="send_state"):
                    if json_data['isroot']:#是根节点
                        if self.root!=json_data['ip']:
                            self.root=json_data['ip']
                        self.root_state=json_data['state']
                    else:
                        if json_data['ip'] in self.down_server:
                            ip_index=self.down_server.index(json_data['ip'])
                            down_server_state[ip_index]=json_data['state']
                        else:
                            self.n_down_server+=1
                            self.down_server.append(json_data['ip'])
                            self.down_server_state.append(json_data['state'])
            except socket.timeout:
                pass
        return
    def udp_send(self,address,message):
        client_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        client_socket.sendto(message, address) #将msg内容发送给指定接收方
        client_socket.close();
        return
    def show_state(self):
        while self.running:
            print('\r','root:{},state:{},down_server:{},state:{}'.format(self.root,self.root_state,self.down_server,self.down_server_state),end='')
            time.sleep(0.5)
        return
    def time_task(self):
        while self.running:
            time.sleep(60*5)
            print("timed task: start FL")
            self.udp_send((self.root,self.port),self.sendjson1.encode('utf-8'))
        return
    def start(self):
        for i in self.order:
            print(i,self.order[i])
        self.th_listen= threading.Thread(target=self.listen, args=())
        self.th_listen.start()
        self.th_show= threading.Thread(target=self.show_state, args=())
        self.th_show.start()
        self.th_time_task= threading.Thread(target=self.time_task, args=())
        self.th_time_task.start()
        while self.running:
            a = input()
            if a=="1":
                self.udp_send((self.root,self.port),self.sendjson1.encode('utf-8'))
            elif a=="0":
                self.running=False
        return



c=FLcloud()
c.start()
